// QUIRKSIM
// Copyright (c) 2009 Gregory Andrew Robinson (garobinson@ucdavis.edu)
//
// Permission is hereby granted, free of charge, to any person
// obtaining a copy of this software and associated documentation
// files (the "Software"), to deal in the Software without
// restriction, including without limitation the rights to use,
// copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the
// Software is furnished to do so, subject to the following
// conditions:
//
// The above copyright notice and this permission notice shall be
// included in all copies or substantial portions of the Software.
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
// EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
// OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
// HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
// WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
// OTHER DEALINGS IN THE SOFTWARE.

#ifndef QUIRKSIM_CC_
#define QUIRKSIM_CC_

#include <stdexcept>
#include <gflags/gflags.h>
#include "quirksim.h"

// Implementing the undefined functions which Hep3Vector is _supposed_ to have.
// It is pretty absurd that I have to do this - is something wrong here?
// TODO(grobinson): Send these upstream if they really aren't in the codebase.
namespace CLHEP {

bool Hep3Vector::isNear(const Hep3Vector& v2, double epsilon) const {
  Hep3Vector diff = *this - v2;
  if (diff.mag() < epsilon) {
    return true;
  } else {
    return false;
  }
}

Hep3Vector Hep3Vector::project(const Hep3Vector& v2) const {
  return v2.unit() * this->dot(v2);
}
}  // namespace CLHEP

XPPoint::XPPoint(const Hep3Vector& x, const Hep3Vector& p)
  : x_(x),
    p_(p) {}
XPPoint::~XPPoint() {}

ParamPoint::ParamPoint(const Hep3Vector& x, const Hep3Vector& p)
  : XPPoint(x, p) {}
ParamPoint::~ParamPoint() {}

Quirk::Quirk(const Hep3Vector& x, const Hep3Vector& p, double mass,
             const QuirkSystem& sys, const bool lambda)
  : ParamPoint(x, p),
    mass_(mass),
    sys_(sys),
    lambda_(lambda) {
  RegisterForce(new QuirkForce(*this));
}

Quirk::~Quirk() {}

ParamPoint* Quirk::Clone(const QuirkSystem& to_sys) const {
  ParamPoint* to = new Quirk(x(), p(), mass(), to_sys, lambda_);
  return to;
}

void Quirk::ComputeDerivs(const QuirkSystem& in_state,
                          String_T::const_iterator it,
                          XPDeriv* out) const {
  out->dp_dt = DpDt();
  out->dx_dt = DxDt();
}

Hep3Vector Quirk::DxDt() const {
  return 1.0 / sqrt(mass_ * mass_ + p() * p()) * p();
}

Hep3Vector Quirk::OutwardUhat() const {
  Hep3Vector out;
  if(lambda_ == 0) {
    out = sys_.Uhat(sys_.points().begin());
    out *= -1.0;
    return out;
  } else {
    out = sys_.Uhat(sys_.points_->end() - 1);
    return out;
  }
}

Hep3Vector Quirk::VPar() const {
  Hep3Vector v = DxDt();
  Hep3Vector out_Uhat = OutwardUhat();
  return DxDt().project(OutwardUhat());
}

Hep3Vector Quirk::VPerp() const {
  return DxDt().perpPart(OutwardUhat());
}

Hep3Vector Quirk::DpDt() const {
  Hep3Vector netforce = Hep3Vector();
  std::vector<ForceInterface*>::const_iterator f_iter = forces_.begin();
  for (f_iter = forces_.begin();
       f_iter != forces_.end();
       f_iter++) {
    netforce += (*f_iter)->Force();
  }
  return netforce;
}

void Quirk::RegisterForce(ForceInterface* f) {
  forces_.push_back(f);
}

FixedEnd::FixedEnd(const Hep3Vector& x)
  : ParamPoint(x, Hep3Vector(0, 0, 0)){

}

ParamPoint* FixedEnd::Clone(const QuirkSystem& sys) const {
  return new FixedEnd(x());
}

void FixedEnd::ComputeDerivs(const QuirkSystem& in_state,
                             String_T::const_iterator it,
                             XPDeriv* out) const {
  out->dp_dt = 0;
  out->dx_dt = 0;
}

StringPoint::StringPoint(const Hep3Vector& x,
                         const Hep3Vector& p,
                         const QuirkSystem& sys)
  : ParamPoint(x, p),
    sys_(sys) {
}

StringPoint::~StringPoint() {}

ParamPoint* StringPoint::Clone(const QuirkSystem& to_sys) const {
  ParamPoint* to = new StringPoint(x(), p(), to_sys);
  return to;
}

void StringPoint::ComputeDerivs(const QuirkSystem& in_state,
                                String_T::const_iterator it,
                                XPDeriv* out) const {
  out->dp_dt = DpDt(it);
  out->dx_dt = DxDt();
}

Hep3Vector StringPoint::DxDt() const {
  return p() * sqrt(1 / (1 + p() * p()));
}

void Print(const Hep3Vector& v) {
  printf("%.20g, %.20g, %.20g", v.x(), v.y(), v.z());
}

Hep3Vector StringPoint::DpDt(String_T::const_iterator it) const {
  const ParamPoint& point_a(**(it - 1));
  const ParamPoint& point_b(**(it + 1));

  Hep3Vector Pi_a = point_a.Pi();
  Hep3Vector Pi_b = point_b.Pi();

  Hep3Vector u_a = sys_.Uhat(it - 1);
  Hep3Vector u_b = sys_.Uhat(it + 1);

  Hep3Vector mu_a = 1.0 / sqrt(1 + Pi_a.mag2()) * u_a;
  Hep3Vector mu_b = 1.0 / sqrt(1 + Pi_b.mag2()) * u_b;

  double dl = sys_.DeltaL(it - 1, it + 1);

  Hep3Vector dmu_dl = (mu_b - mu_a) * (1.0 / dl);
  Hep3Vector dPi_dl = (Pi_b - Pi_a) * (1.0 / dl);

  Hep3Vector pi = Pi();

  return dmu_dl - pi * (1 / sqrt(1 + pi.mag2())) * (sys_.Uhat(it).dot(dPi_dl));
}

QuirkSystem::QuirkSystem(double sigma, double point_separation)
  : sigma_(sigma), separation_(point_separation) {
  points_ = new std::vector<ParamPoint*>();
}

QuirkSystem::~QuirkSystem() {
  for (std::vector<ParamPoint*>::iterator it = points_->begin();
       it < points_->end();
       it++) {
    delete *it;
  }
  delete points_;
}

QuirkSystem* QuirkSystem::Clone() const {
  QuirkSystem* sys = new QuirkSystem(sigma_, separation_);
  for(std::vector<ParamPoint*>::const_iterator from = this->points_->begin();
      from < points_->end();
      from++) {
    ParamPoint* point_clone = (*from)->Clone(*sys);
    sys->mutable_points()->push_back(point_clone);
  }
  return sys;
}

void QuirkSystem::Add(const std::vector<XPDeriv>& add, double factor) {
  if (add.size() != points_->size()) {
    throw std::length_error("input size does not match points_ size.");
  }

  for (unsigned int i = 0; i < points_->size(); i++) {
    points_->at(i)->Add(add.at(i), factor);
  }
}

void QuirkSystem::ComputeDerivs(const QuirkSystem& in_state,
                                std::vector<XPDeriv>* out) {
  out->resize(in_state.points_->size());
  std::vector<XPDeriv>::iterator out_it = out->begin();
  for (String_T::const_iterator it = in_state.points().begin();
       it != in_state.points().end();
       it++) {
    (*it)->ComputeDerivs(in_state, it, &(*out_it));
    ++out_it;
  }
}

double QuirkSystem::DeltaL(String_T::const_iterator min,
                           String_T::const_iterator max) const {
  double dl = 0;
  for (; min != max; ++min) {
    dl += ((*(min + 1))->x() - (*min)->x()).mag();
  }
  return dl;
}

double QuirkSystem::Length() const {
  double length;
  Hep3Vector previous = points_->at(0)->x();
  for (unsigned int i = 1; i < points_->size(); i++) {
    length += (points_->at(i)->x() - previous).mag();
    previous = points_->at(i)->x();
  }
  return length;
}

Hep3Vector QuirkSystem::Uhat(String_T::const_iterator point) const {
  if (point == points_->begin()) {
    return ((*(point + 1))->x() - (*point)->x()).unit();
  } else if (point == points_->end() - 1) {
    return ((*point)->x() - (*(point - 1))->x()).unit();
  } else {
    // Hep3Vector dx_0 = points_->at(index)->x() - points_->at(index - 1)->x();
    // Hep3Vector dx_1 = points_->at(index + 1)->x() - points_->at(index)->x();
    // return (dx_0 * (1 / dx_0.mag2()) + dx_1 * (1 / dx_1.mag2())).unit();
    return ((*(point + 1))->x() - (*(point - 1))->x()).unit();
  }
}

void QuirkSystem::RK4Step(const double dt) {
  std::vector<XPDeriv> A;
  std::vector<XPDeriv> B;
  std::vector<XPDeriv> C;
  std::vector<XPDeriv> D;

  ComputeDerivs(*this, &A);
  QuirkSystem* b_state = this->Clone();
  b_state->Add(A, dt / 2.0);
  //b_state->Reparam(separation_);
  b_state->ConstrainP();

  ComputeDerivs(*b_state, &B);
  QuirkSystem* c_state = this->Clone();
  c_state->Add(B, dt / 2.0);
  //c_state->Reparam(separation_);
  c_state->ConstrainP();

  ComputeDerivs(*b_state, &C);
  QuirkSystem* d_state = this->Clone();
  d_state->Add(C, dt);
  //d_state->Reparam(separation_);
  d_state->ConstrainP();

  ComputeDerivs(*d_state, &D);
  Add(A, dt / 6.0);
  Add(B, dt * 2.0 / 6.0);
  Add(C, dt * 2.0 / 6.0);
  Add(D, dt / 6.0);
  Reparam(separation_);
  ConstrainP();

  delete b_state;
  delete c_state;
  delete d_state;
}

void QuirkSystem::Reparam(double separation) {
  if (separation < 0) return;
  int i = 0;
  for (String_T::iterator point = points_->begin();
       point != points_->end() - 2;
       ) {  //point is incremented within succeeding conditional.
    double dl = DeltaL(point, point+1);
    if (dl > 2 * separation) {
      Hep3Vector new_x((*(point+1))->x() + (*point)->x());
      new_x /= 2;
      Hep3Vector new_p((*(point+1))->p() + (*point)->p());
      new_p /= 2;
      // insert operation may invalidate iterator.
      point = points_->insert(point + 1, new StringPoint(new_x, new_p, *this));
      // Since we set point to the inserted point, we must go back to the
      // current position at hand.
      --point;
    } else if (dl < separation / 2) {
      if (point + 1 == --points_->end()) {
        // Delete the current point, which is right next to the last.
        STLDeleteContainerPointers(point, --points_->end());
        points_->erase(point, --points_->end());
        --point;
        --i;
      } else {
        // Delete next point, since the last span is fine.
        STLDeleteContainerPointers(point + 1, point + 2);
        points_->erase(point + 1, point + 2);
      }
    } else {
      ++point;
      ++i;
    }
  }
}

QuirkForce::QuirkForce(const Quirk& quirk)
  : quirk_(quirk) {
}

Hep3Vector QuirkForce::Force() const {
  const Hep3Vector outward_normal = quirk_.OutwardUhat();
  const Hep3Vector v_perp = quirk_.VPerp();
  const Hep3Vector v_par = quirk_.VPar();
  return -1 * quirk_.sys().sigma() * (
      sqrt(1 - v_perp * v_perp) * outward_normal.unit()
      + v_par * v_par / sqrt(1 - v_perp * v_perp) * v_perp);
}

#endif  // QUIRKSIM_CC_
